import sqlite3
import os
import time
import base64
import zlib
import argparse
import glob

class FileDBManager:
    def process_directory(self, path, base_dir=None):
        """
        递归处理目录结构
        :param path: 目录路径
        :param base_dir: 基准目录（用于生成相对路径）
        """
        for root, dirs, files in os.walk(path):
            for file in files:
                full_path = os.path.join(root, file)
                self.add_file(full_path, base_dir or path)

    def batch_add_files(self, file_paths, recursive=False):
        """
        批量添加文件（支持目录和通配符）
        :param file_paths: 文件/目录路径列表
        :param recursive: 是否递归处理目录
        """
        for path in file_paths:
            if os.path.isdir(path):
                if recursive:
                    self.process_directory(path)
                else:
                    print(f"跳过目录 {path}（未启用递归选项）")
            else:
                for file in glob.glob(path, recursive=recursive):
                    if os.path.isfile(file):
                        self.add_file(file)

    def batch_extract_files(self, filenames, output_dir=None):
        """
        批量从数据库提取文件
        :param filenames: 文件名列表
        """
        for name in filenames:
            try:
                self.extract_file(name)
            except Exception as e:
                print(f"提取文件 {name} 时出错: {str(e)}")
    
    def __init__(self, db_name='files.db'):
        self.db_name = db_name
        self.create_table()

    def create_table(self):
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute('''CREATE TABLE IF NOT EXISTS files (
            encrypted_filename TEXT PRIMARY KEY,
            content BLOB NOT NULL,
            relative_path TEXT NOT NULL,
            file_size INTEGER NOT NULL,
            modified_time REAL NOT NULL
        )''')
        conn.commit()
        conn.close()

    def _encrypt_filename(self, filename):
        compressed = zlib.compress(filename.encode())
        encrypted = base64.b64encode(compressed).decode()
        return encrypted

    def _decrypt_filename(self, encrypted_filename):
        compressed = base64.b64decode(encrypted_filename)
        filename = zlib.decompress(compressed).decode()
        return filename

    def handle_duplicate_filename(self, filename):
        base_name, extension = os.path.splitext(filename)
        counter = 1
        new_filename = filename
        while True:
            encrypted_filename = self._encrypt_filename(new_filename)
            conn = sqlite3.connect(self.db_name)
            c = conn.cursor()
            c.execute("SELECT EXISTS(SELECT 1 FROM files WHERE encrypted_filename=?)", (encrypted_filename,))
            result = c.fetchone()[0]
            conn.close()
            if not result:
                break
            new_filename = f"{base_name}_{counter}{extension}"
            counter += 1
        return new_filename

    def add_file(self, file_path, base_dir=None):
        file_stat = os.stat(file_path)
        try:
            with open(file_path, 'rb') as file:
                filedata = file.read()
            if base_dir:
                relative_path = os.path.relpath(file_path, base_dir)
            else:
                relative_path = os.path.basename(file_path)
            filename = os.path.basename(file_path)
            encrypted_filename = self._encrypt_filename(filename)
            
            # 获取文件元数据
            file_size = file_stat.st_size
            modified_time = file_stat.st_mtime
            conn = sqlite3.connect(self.db_name)
            c = conn.cursor()
            c.execute("SELECT EXISTS(SELECT 1 FROM files WHERE encrypted_filename=?)", (encrypted_filename,))
            result = c.fetchone()[0]
            if result:
                choice = input(f"文件 {filename} 已存在于数据库中，是否自动重命名保存？(y/n): ")
                if choice.lower() == 'y':
                    new_filename = self.handle_duplicate_filename(filename)
                    encrypted_filename = self._encrypt_filename(new_filename)
                else:
                    print("操作取消。")
                    return
            # add_time = time.time()
            c.execute("""INSERT INTO files VALUES (?,?,?,?,?)
                          ON CONFLICT(encrypted_filename) DO UPDATE SET
                          content=excluded.content,
                          file_size=excluded.file_size,
                          modified_time=excluded.modified_time""",
                         (encrypted_filename, filedata, relative_path, file_size, modified_time))
            conn.commit()
            conn.close()
            print(f"文件 {filename} 已成功添加到数据库。")
        except FileNotFoundError:
            print(f"错误: 文件 {file_path} 未找到。")

    def extract_file(self, filename, output_dir=None):
        """
        从数据库提取文件到指定目录
        :param filename: 文件名
        :param output_dir: 输出目录（默认为原路径）
        """
        encrypted_filename = self._encrypt_filename(filename)
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute("SELECT content FROM files WHERE encrypted_filename=?", (encrypted_filename,))
        result = c.fetchone()
        conn.close()
        
        if not result:
            raise FileNotFoundError(f"文件 {filename} 未在数据库中找到")
            
        target_path = os.path.join(output_dir or os.getcwd(), filename)
        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        
        with open(target_path, 'wb') as file:
            file.write(result[0])
        print(f"文件已提取到: {target_path}")

    def extract_all_files(self, output_dir):
        """
        提取数据库中的所有文件到指定目录
        :param output_dir: 输出目录
        """
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute('SELECT encrypted_filename, content FROM files')
        
        os.makedirs(output_dir, exist_ok=True)
        count = 0
        
        for encrypted_filename, filedata in c.fetchall():
            try:
                filename = self._decrypt_filename(encrypted_filename)
                target_path = os.path.join(output_dir, filename)
                with open(target_path, 'wb') as f:
                    f.write(filedata)
                count += 1
            except Exception as e:
                print(f"提取文件 {filename} 时出错: {str(e)}")
        
        print(f"成功提取 {count} 个文件到目录: {output_dir}")
        conn.close()

    def query_file(self, filename):
        encrypted_filename = self._encrypt_filename(filename)
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute("SELECT EXISTS(SELECT 1 FROM files WHERE encrypted_filename=?)", (encrypted_filename,))
        result = c.fetchone()[0]
        conn.close()
        if result:
            print(f"文件 {filename} 存在于数据库中。")
        else:
            print(f"文件 {filename} 不存在于数据库中。")

    def delete_file(self, filename):
        encrypted_filename = self._encrypt_filename(filename)
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute("DELETE FROM files WHERE encrypted_filename=?", (encrypted_filename,))
        if c.rowcount > 0:
            conn.commit()
            print(f"文件 {filename} 已成功从数据库中删除。")
        else:
            print(f"错误: 文件 {filename} 未在数据库中找到。")
        conn.close()

    def update_file(self, file_path):
        try:
            with open(file_path, 'rb') as file:
                filedata = file.read()
            filename = os.path.basename(file_path)
            encrypted_filename = self._encrypt_filename(filename)
            conn = sqlite3.connect(self.db_name)
            c = conn.cursor()
            c.execute("UPDATE files SET filedata=?, add_time=? WHERE encrypted_filename=?",
                      (filedata, time.time(), encrypted_filename))
            if c.rowcount > 0:
                conn.commit()
                print(f"文件 {filename} 已成功在数据库中更新。")
            else:
                print(f"错误: 文件 {filename} 未在数据库中找到。")
            conn.close()
        except FileNotFoundError:
            print(f"错误: 文件 {file_path} 未找到。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='文件数据库管理工具')
    subparsers = parser.add_subparsers(dest='command', required=True)

    # 添加命令
    add_parser = subparsers.add_parser('add', help='添加文件到数据库')
    add_parser.add_argument('paths', nargs='+', help='文件/目录路径（支持通配符）')
    add_parser.add_argument('-r', '--recursive', action='store_true', help='递归处理目录')
    add_parser.add_argument('-db', '--database', default='files.db', help='数据库文件名')

    # 提取命令
    extract_parser = subparsers.add_parser('extract', help='从数据库提取文件')
    extract_parser.add_argument('files', nargs='*', help='要提取的文件名（支持通配符）')
    extract_parser.add_argument('-A', '--all', action='store_true', help='提取所有文件')
    extract_parser.add_argument('-o', '--output', help='输出目录')
    extract_parser.add_argument('-db', '--database', default='files.db', help='数据库文件名')

    # 查询命令
    query_parser = subparsers.add_parser('query', help='查询数据库信息')
    query_parser.add_argument('pattern', nargs='?', default='*', help='查询模式（支持通配符）')
    query_parser.add_argument('-db', '--database', default='files.db', help='数据库文件名')

    # 删除命令
    del_parser = subparsers.add_parser('delete', help='删除数据库文件')
    del_parser.add_argument('files', nargs='+', help='要删除的文件名（支持通配符）')
    del_parser.add_argument('-db', '--database', default='files.db', help='数据库文件名')

    # 更新命令
    update_parser = subparsers.add_parser('update', help='更新已修改的文件')
    update_parser.add_argument('paths', nargs='+', help='要检查的文件/目录路径')
    update_parser.add_argument('-r', '--recursive', action='store_true', help='递归检查目录')
    update_parser.add_argument('-db', '--database', default='files.db', help='数据库文件名')

    args = parser.parse_args()

    manager = FileDBManager(args.database)

    if args.command == 'add':
        manager.batch_add_files(args.paths, recursive=args.recursive)
    elif args.command == 'extract':
        if args.all:
            manager.extract_all_files(args.output or os.getcwd())
        else:
            manager.batch_extract_files(args.files, args.output or os.getcwd())
    elif args.command == 'query':
        conn = sqlite3.connect(manager.db_name)
        c = conn.cursor()
        c.execute("SELECT encrypted_filename FROM files WHERE relative_path LIKE ?", (args.pattern.replace('*','%'),))
        for (encrypted_name,) in c.fetchall():
            print(manager._decrypt_filename(encrypted_name))
        conn.close()
    elif args.command == 'delete':
        for pattern in args.files:
            for file in glob.glob(pattern):
                manager.delete_file(os.path.basename(file))
    elif args.command == 'update':
        # 实现更新逻辑
        pass
    